# -*- coding: utf-8 -*-

import shapefile


def save_shp(target, shape_type, geo_data, field_names, field_types, field_sizes, field_decimals, field_values):
    """
    target => 生成的文件名, 自动生成.shp .dbf .shx, 如"D:\\test"。
    shapeType => 几何类型,说明如下。
        NULL = 0, POINT = 1, POLYLINE = 3, POLYGON = 5, MULTIPOINT = 8
        POINTZ = 11, POLYLINEZ = 13, POLYGONZ = 15, MULTIPOINTZ = 18
        POINTM = 21, POLYLINEM = 23, POLYGONM = 25, MULTIPOINTM = 28, MULTIPATCH = 31
    geo_data => 几何数据
        points = [[1,1],[2,2]]
        polylines = [[[[1,1],[2,2]],[[2,2],[3,3]]],[[[1,2],[2,4]]]]
        polygons = [[[[1,1],[2,2],[3,1]],[[3,1],[2,2],[3,2]]],[[[1,2],[2,3],[1,3]]]]
        pointzs = [[1,1,1],[2,2,2]]
        polylinezs = [[[[1,1,1],[2,2,2]],[[2,2,3],[3,3,4]]],[[[1,2,5],[2,4,6]]]]
        polygonzs = [[[[1,1,1],[2,2,2],[3,1,3]],[[3,1,4],[2,2,5],[3,2,6]]],[[[1,2,7],[2,3,8],[1,3,9]]]]
    field_names => 属性名称,如: ['test1','test2','test3']
    field_types => 属性类型, 如: ['C','N','F']
        “C”: 字符,文字。
        “N”: 数字,带或不带小数。
        “F”: 浮动(与“N”相同)。
        “L”: 逻辑,表示布尔值True、False值。
        “D”: 日期。
        “M”: 备忘录,在GIS中没有意义,而是xbase规范的一部分。
    field_sizes => 段长度,默认50, 如; [50, 50, 50]  
    field_decimals => 字段精度,默认为0, 如: [0, 0, 2]
    field_values => 属性值, 如: [['test', 1, 1.111], ['测试', 2, 2.226]]
    """
    w = shapefile.Writer(target=target, shapeType=shape_type, autoBalance=1)
        
    # 添加属性
    for field in range(len(field_names)):
        w.field(field_names[field], field_types[field], field_sizes[field], field_decimals[field])
        
    # 写出几何数据
    for record in range(len(geo_data)):
        if shape_type == 1:
            w.point(*geo_data[record])
        elif shape_type == 3:
            w.line(geo_data[record])
        elif shape_type == 5:
            w.poly(geo_data[record])
        elif shape_type == 11:
            w.pointz(*geo_data[record])
        elif shape_type == 13:
            w.linez(geo_data[record])
        elif shape_type == 15:
            w.polyz(geo_data[record])
        else:
            continue

        # 写出属性数据
        records = []
        for field in range(len(field_names)):
            records.append(field_values[record][field])
        w.record(*records)


def get_shp_fields(path):
    try:
        file = shapefile.Reader(path)
    except UnicodeDecodeError:
        file = shapefile.Reader(path, encoding="gbk")

    fields, result = file.fields, []
    for field in fields:
        result.append(field[0])
    
    return result


def get_shp_shapes(path, fields):
    try:
        file = shapefile.Reader(path)
        shape_records = file.shapeRecords()
    except UnicodeDecodeError:
        file = shapefile.Reader(path, encoding="gbk")
        shape_records = file.shapeRecords()
                
    shp_types, shp_datas, shp_fields = [], [], []
    for _ in range(len(fields)):
        field_values = []
        shp_fields.append(field_values)

    for shape_record in shape_records:
        shp_type = shape_record.shape.shapeType
        if shp_type == 1:
            shp_types.append('point')
            points = shape_record.shape.points
            points_order = []
            points_order.append(len(points))
            for point in points:
                points_order.append(point[0])
                points_order.append(point[1])
            shp_datas.append(points_order)
        elif shp_type == 3:
            shp_types.append('polyline')
            points = shape_record.shape.points
            parts = shape_record.shape.parts
            polyline_record = []
            for idx in range(0,len(parts)-1):
                polyline_record_part = points[parts[idx]:parts[idx+1]]
                polyline_record.append(polyline_record_part)
            polyline_record_part = points[parts[-1]:]
            polyline_record.append(polyline_record_part)
            shp_datas.append(polyline_record)
        elif shp_type == 5:
            shp_types.append('polygon')
            points = shape_record.shape.points
            parts = shape_record.shape.parts
            polygon_record = []
            for idx in range(0,len(parts)-1):
                polygon_record_part = points[parts[idx]:parts[idx+1]]
                polygon_record.append(polygon_record_part)
            polygon_record_part = points[parts[-1]:]
            polygon_record.append(polygon_record_part)
            shp_datas.append(polygon_record)
        elif shp_type == 11:
            shp_types.append('pointz')
            points = shape_record.shape.points
            zs = shape_record.shape.z
            points_order = []
            points_order.append(len(points))
            for idx in range(len(points)):
                points_order.append(points[idx][0])
                points_order.append(points[idx][1])
                points_order.append(zs[idx])
            shp_datas.append(points_order)
        elif shp_type == 13:
            shp_types.append('polylinez')
            points = shape_record.shape.points
            zs = shape_record.shape.z
            parts = shape_record.shape.parts
            polylinez_record = []
            for idx in range(0,len(parts)-1):
                polylinez_record_part = []
                for idx_pt in range(parts[idx],parts[idx+1]):
                    pt = list(points[idx_pt])
                    pt.append(zs[idx_pt])
                    polylinez_record_part.append(pt)
                polylinez_record.append(polylinez_record_part)
            polylinez_record_part = []
            for idx_pt in range(parts[-1],len(points)):
                pt = list(points[idx_pt])
                pt.append(zs[idx_pt])
                polylinez_record_part.append(pt)
            polylinez_record.append(polylinez_record_part)
            shp_datas.append(polylinez_record)
        elif shp_type == 15:
            shp_types.append('polygonz')
            points = shape_record.shape.points
            zs = shape_record.shape.z
            parts = shape_record.shape.parts
            polygonz_record = []
            for idx in range(0,len(parts)-1):
                polygonz_record_part = []
                for idx_pt in range(parts[idx],parts[idx+1]):
                    pt = list(points[idx_pt])
                    pt.append(zs[idx_pt])
                    polygonz_record_part.append(pt)
                polygonz_record.append(polygonz_record_part)
            polygonz_record_part = []
            for idx_pt in range(parts[-1],len(points)):
                pt = list(points[idx_pt])
                pt.append(zs[idx_pt])
                polygonz_record_part.append(pt)
            polygonz_record.append(polygonz_record_part)
            shp_datas.append(polygonz_record)
        else:
            print("undefined type")

        for idx, field in enumerate(fields):
            shp_fields[idx].append(shape_record.record[field])

        return shp_types, shp_datas, shp_fields


def add_shp_records(path, geo_data, field_values):
    # 先读取原文件的几何数据和属性数据
    try:
        file = shapefile.Reader(path)
        shape_records = file.shapeRecords()
        shape_fields = file.fields
        shape_type = file.shapeType
    except UnicodeDecodeError:
        file = shapefile.Reader(path, encoding="gbk")
        shape_records = file.shapeRecords()
        shape_fields = file.fields
        shape_type = file.shapeType   

    # 创建与原文件同名的导出文件
    w = shapefile.Writer(target=path, shapeType=shape_type, autoBalance=1)
        
    # 创建与原文件相同的属性
    for shape_field in shape_fields:
        w.field(*shape_field)
        
    # 将原文件的几何数据和属性数据添加到新文件
    for shape_record in shape_records:
        w.shape(shape_record.shape)
        w.record(*shape_record.record)

    # 在新文件追加待添加的几何数据和属性数据
    for idx_record in range(len(geo_data)):
        # 追加几何数据
        if shape_type == 1:
            w.point(*geo_data[idx_record])
        elif shape_type == 3:
            w.line(geo_data[idx_record])
        elif shape_type == 5:
            w.poly(geo_data[idx_record])
        elif shape_type == 11:
            w.pointz(*geo_data[idx_record])
        elif shape_type == 13:
            w.linez(geo_data[idx_record])
        elif shape_type == 15:
            w.polyz(geo_data[idx_record])
        else:
            continue

        # 追加属性数据
        w.record(*field_values[idx_record])


# 测试数据
if __name__ == '__main__':
    field_names = ['test1', 'test2', 'test3']
    field_types = ['C','N', 'F']
    field_sizes = [50, 50, 50]
    field_decimals = [0, 0, 2]
    field_values = [['test', 1, 1.111], ['测试', 2, 2.226]]

    target = "point"
    shape_type = 1
    points = [[1, 1], [2, 2]]
    print(save_shp(target, shape_type, points, field_names, field_types, field_sizes, field_decimals, field_values))

    target = "polyline"
    shape_type = 3
    polylines = [[[[1, 1], [2, 2]], [[2, 2], [3, 3]]], [[[1, 2], [2, 4]]]]
    print(save_shp(target, shape_type, polylines, field_names, field_types, field_sizes, field_decimals, field_values))

    target = "polygon"
    shape_type = 5
    polygons = [[[[1, 1], [2, 2], [3, 1]], [[3, 1], [2, 2], [3, 2]]], [[[1, 2], [2, 3], [1, 3]]]]
    print(save_shp(target, shape_type, polygons, field_names, field_types, field_sizes, field_decimals, field_values))

    target = "pointz"
    shape_type = 11
    pointzs = [[1, 1, 1], [2, 2, 2]]
    print(save_shp(target, shape_type, pointzs, field_names, field_types, field_sizes, field_decimals, field_values))

    target = "polylinez"
    shape_type = 13
    polylinezs = [[[[1, 1, 1], [2, 2, 2]], [[2, 2, 3], [3, 3, 4]]], [[[1, 2, 5], [2, 4, 6]]]]
    print(save_shp(target, shape_type, polylinezs, field_names, field_types, field_sizes, field_decimals, field_values))

    target = "polygonz"
    shape_type = 15
    polygonzs = [[[[1, 1, 1], [2, 2, 2], [3, 1, 3]], [[3, 1, 4], [2, 2, 5], [3, 2, 6]]], [[[1, 2, 7], [2, 3, 8], [1, 3, 9]]]]
    print(save_shp(target, shape_type, polygonzs, field_names, field_types, field_sizes, field_decimals, field_values))

    print(get_shp_fields("polygon"))
    print(get_shp_shapes("point", ['test1', 'test3']))
    print(get_shp_shapes("polyline", ['test1', 'test3']))
    print(get_shp_shapes("polygon", ['test1', 'test3']))
    print(get_shp_shapes("pointz", ['test1', 'test3']))
    print(get_shp_shapes("polylinez", ['test1', 'test3']))
    print(get_shp_shapes("polygonz", ['test1', 'test3']))

    path = 'point'
    points = [[3, 3], [4, 4]]
    field_values = [['add', 3, None], ['增加', 4, None]]
    print(add_shp_records(path, points, field_values))
